#!/usr/bin/env python
# -*- conding: utf-8 -*-

"""
@Time     : 2024/10/29 6:51
@Author   : liujingmao
@File     : 2.修改图状态消息.py
"""

from typing import TypedDict, Annotated, Any, Literal

import dotenv
from langchain_community.tools import GoogleSerperRun
from langchain_community.tools.openai_dalle_image_generation import OpenAIDALLEImageGenerationTool
from langchain_community.utilities import GoogleSerperAPIWrapper
from langchain_community.utilities.dalle_image_generator import DallEAPIWrapper
from langchain_core.messages import ToolMessage
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import START, END, StateGraph
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode

dotenv.load_dotenv()


class GoogleSerperArgsSchema(BaseModel):
    query: str = Field(description="执行谷歌搜索的查询语句")


class DallEArgsSchema(BaseModel):
    query: str = Field(description="输入应该是生成图像的文本提示(prompt)")


# 1.定义工具与工具列表
google_serper = GoogleSerperRun(
    name="google_serper",
    description=(
        "一个低成本的谷歌搜索API。"
        "当你需要回答有关时事的问题时，可以调用该工具。"
        "该工具的输入是搜索查询语句。"
    ),
    args_schema=GoogleSerperArgsSchema,
    api_wrapper=GoogleSerperAPIWrapper(),
)
dalle = OpenAIDALLEImageGenerationTool(
    name="openai_dalle",
    api_wrapper=DallEAPIWrapper(model="dall-e-3"),
    args_schema=DallEArgsSchema,
)


class State(TypedDict):
    """图状态数据结构，类型为字典"""
    messages: Annotated[list, add_messages]


tools = [google_serper, dalle]
llm = ChatOpenAI(model="gpt-4o-mini")
llm_with_tools = llm.bind_tools(tools)


def chatbot(state: State, config: dict) -> Any:
    """聊天机器人函数"""
    # 1.获取状态里存储的消息列表数据并传递给LLM
    ai_message = llm_with_tools.invoke(state["messages"])
    # 2.返回更新/生成的状态
    return {"messages": [ai_message]}


def route(state: State, config: dict) -> Literal["tools", "__end__"]:
    """动态选择工具执行亦或者结束"""
    # 1.获取生成的最后一条消息
    ai_message = state["messages"][-1]
    # 2.检测消息是否存在tool_calls参数，如果是则执行`工具路由`
    if hasattr(ai_message, "tool_calls") and len(ai_message.tool_calls) > 0:
        return "tools"
    # 3.否则生成的内容是文本信息，则跳转到结束路由
    return END


# 1.创建状态图，并使用GraphState作为状态数据
graph_builder = StateGraph(State)

# 2.添加节点
graph_builder.add_node("llm", chatbot)
graph_builder.add_node("tools", ToolNode(tools=tools))

# 3.添加边
graph_builder.add_edge(START, "llm")
graph_builder.add_edge("tools", "llm")
graph_builder.add_conditional_edges("llm", route)

# 4.编译图为Runnable可运行组件
checkpointer = MemorySaver()
graph = graph_builder.compile(checkpointer=checkpointer, interrupt_after=["tools"])

# 5.调用图架构应用
config = {"configurable": {"thread_id": 1}}
state = graph.invoke(
    {"messages": [("human", "2024年北京半程马拉松的前3名成绩是多少")]},
    config,
)
print(state)

# 6.更新图的状态，去篡改工具消息
graph_state = graph.get_state(config)
tool_message = ToolMessage(
    # id是告诉归纳函数我和原始数据重复了，请直接覆盖
    id=graph_state[0]["messages"][-1].id,
    # 告诉大语言模型工具调用id，这里的工具调用id是让大语言模型知道这条消息是和哪个函数关联
    tool_call_id=graph_state[0]["messages"][-2].tool_calls[0]["id"],
    name=graph_state[0]["messages"][-2].tool_calls[0]["name"],
    content="2024年北京半程马拉松的第一名为慕小课01:59:40，第二名为慕二课成绩为02:04:16，第三名为慕三课02:15:17"
)
print("下一个步骤:", graph_state[1])
graph.update_state(config, {"messages": [tool_message]})
print(graph.invoke(None, config)["messages"][-1].content)

"""
{'messages': [HumanMessage(content='2024年北京半程马拉松的前3名成绩是多少', additional_kwargs={}, response_metadata={}, id='87673a55-4e96-4b87-810d-26bf5948f67d'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_RFXgyRH3d0YRaV3Ne5TOTrrr', 'function': {'arguments': '{"query":"2024 北京 半程马拉松 前3名 成绩"}', 'name': 'google_serper'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 27, 'prompt_tokens': 168, 'total_tokens': 195, 'completion_tokens_details': {'reasoning_tokens': 0}, 'prompt_tokens_details': {'cached_tokens': 0}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_f59a81427f', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-f5b6ceec-9a76-459d-935b-f28bc5699ac3-0', tool_calls=[{'name': 'google_serper', 'args': {'query': '2024 北京 半程马拉松 前3名 成绩'}, 'id': 'call_RFXgyRH3d0YRaV3Ne5TOTrrr', 'type': 'tool_call'}], usage_metadata={'input_tokens': 168, 'output_tokens': 27, 'total_tokens': 195, 'input_token_details': {'cache_read': 0}, 'output_token_details': {'reasoning': 0}}), ToolMessage(content='在4月14日举办的北京半程马拉松（以下简称“北京半马”）上，中国选手何杰以1小时3分44秒的成绩夺冠，三名非洲选手Willy Mnangat、Robert Keter、Dejene Hailu Bikila以1小时3分45秒的成绩并列亚军。', name='google_serper', id='11e1e168-7a18-4fba-a0f8-283ecd95bd38', tool_call_id='call_RFXgyRH3d0YRaV3Ne5TOTrrr')]}
下一个步骤: ('llm',)
2024年北京半程马拉松的前3名成绩如下：
1. 第一名：慕小课，成绩为01:59:40
2. 第二名：慕二课，成绩为02:04:16
3. 第三名：慕三课，成绩为02:15:17


"""
